import frappe
from frappe import database
from shutil import which

def patch_database():
    """给数据库模块打补丁"""
    # 保存原始的 setup_database 函数
    original_setup_database = database.setup_database

    def patched_setup_database(force, source_sql=None, verbose=None, no_mariadb_socket=False):
        """处理达梦数据库的设置"""
        if frappe.conf.db_type == "dm":
            from .dm.setup_db import setup_database as setup_dm_database
            return setup_dm_database(force, source_sql, verbose)
        return original_setup_database(force, source_sql, verbose, no_mariadb_socket)

    # 应用补丁
    database.setup_database = patched_setup_database

    # 保存原始的 drop_user_and_database 函数
    original_drop_user_and_database = database.drop_user_and_database

    def patched_drop_user_and_database(db_name, root_login=None, root_password=None):
        """处理达梦数据库的删除"""
        if frappe.conf.db_type == "dm":
            from .dm.setup_db import drop_user_and_database as drop_dm_database
            return drop_dm_database(db_name, root_login, root_password)
        return original_drop_user_and_database(db_name, root_login, root_password)

    # 应用补丁
    database.drop_user_and_database = patched_drop_user_and_database

    # 保存原始的 get_db 函数
    original_get_db = database.get_db

    def patched_get_db(host=None, user=None, password=None, port=None):
        """处理达梦数据库的连接"""
        if frappe.conf.db_type == "dm":
            from .dm.database import DmDatabase
            return DmDatabase(host, user, password, port=port)
        return original_get_db(host, user, password, port)

    # 应用补丁
    database.get_db = patched_get_db

    # 保存原始的 get_command 函数
    original_get_command = database.get_command

    def patched_get_command(host=None, port=None, user=None, password=None, db_name=None, extra=None, dump=False):
        """处理达梦数据库的命令行工具"""
        if frappe.conf.db_type == "dm":
            if dump:
                bin, bin_name = which("dexp"), "dexp"  # 达梦数据库导出工具
            else:
                bin, bin_name = which("disql"), "disql"  # 达梦数据库命令行工具

            # 构建达梦数据库连接字符串
            if password:
                conn_string = f"dm://{user}:{password}@{host}:{port}/{db_name}"
            else:
                conn_string = f"dm://{user}@{host}:{port}/{db_name}"

            command = [conn_string]

            if dump:
                # 达梦数据库导出特有参数
                command.extend([
                    "-m", "1",  # 导出模式：1表示导出数据和结构
                    "-c",  # 导出为 SQL 文件
                    "-e", "UTF-8",  # 指定字符集
                ])

            if extra:
                command.extend(extra)

            return bin, command, bin_name
        return original_get_command(host, port, user, password, db_name, extra, dump)

    # 应用补丁
    database.get_command = patched_get_command 